# calibration.py
# 实现 Bias-Corrected Temperature Scaling (BCTS) 校准器。
# 复用您提供的代码，它非常适合当前任务。

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

class BiasCorrectedTS(nn.Module):
    """
    带偏置修正的温度缩放 (BCTS) 校准器。
    它学习一个全局温度T和每个类别的偏置b。
    """
    def __init__(self, num_classes):
        super(BiasCorrectedTS, self).__init__()
        self.temperature = nn.Parameter(torch.ones(1) * 1.5)
        self.bias = nn.Parameter(torch.zeros(num_classes))

    def forward(self, logits):
        """对输入的logits应用BCTS变换。"""
        temperature = self.temperature.clamp(min=1e-6)  # 只避免除零错误
        return logits / temperature + self.bias

    def fit(self, logits, labels, max_iter=100, lr=0.01):
        """在验证集上学习最优的温度和偏置参数。"""
        self.to(logits.device)
        nll_criterion = nn.CrossEntropyLoss()
        optimizer = optim.LBFGS([self.temperature, self.bias], lr=lr, max_iter=max_iter)

        def eval():
            optimizer.zero_grad()
            calibrated_logits = self(logits)
            loss = nll_criterion(calibrated_logits, labels)
            loss.backward()
            return loss.detach()
        
        optimizer.step(eval)
        print(f"  BCTS fitted. Optimal T: {self.temperature.item():.3f}, Bias norm: {torch.norm(self.bias).item():.3f}")
        return self
